#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
pick_fastest_bun_registry.py
测速并设置 Bun 包管理器默认 registry；默认写入当前目录 ./bunfig.toml，使用 --global 写入 $HOME/.bunfig.toml。
仅用标准库（urllib、ssl、time、argparse、os、shutil、datetime、re）。

术语展开：
- Registry（注册表，软件包索引服务）
- URL = Uniform Resource Locator（统一资源定位符）
- HTTPS = HyperText Transfer Protocol Secure（安全超文本传输协议）
"""

import argparse, os, ssl, time, urllib.request, shutil, re
from datetime import datetime

UA = "bun-registry-bench/1.1 (+https://bun.sh)"
DEFAULT_BYTES = 128 * 1024
DEFAULT_TRIES = 2
DEFAULT_TIMEOUT = 8.0

DEFAULT_CANDIDATES = [
    "https://registry.npmjs.org",
    "https://registry.npmmirror.com",
    "https://mirrors.huaweicloud.com/repository/npm",
    "https://mirrors.cloud.tencent.com/npm",
]

SEARCH_PATH = "/-/v1/search?text=bun&size=1"  # 轻量请求，广泛兼容

def normalize_registry(u: str) -> str:
    u = u.strip()
    if not (u.startswith("http://") or u.startswith("https://")):
        raise ValueError("registry 必须以 http:// 或 https:// 开头")
    return u.rstrip("/") + "/"

def http_read_time(url: str, max_bytes: int, timeout: float):
    ctx = ssl.create_default_context()
    start = time.perf_counter()
    first = None
    read = 0
    try:
        req = urllib.request.Request(url, headers={"User-Agent": UA})
        with urllib.request.urlopen(req, timeout=timeout, context=ctx) as resp:
            chunk = resp.read(1)
            if not chunk:
                return None, "empty"
            first = time.perf_counter()
            read += 1
            while read < max_bytes:
                c = resp.read(min(8192, max_bytes - read))
                if not c:
                    break
                read += len(c)
        elapsed = time.perf_counter() - start
        return elapsed, f"ok rtt≈{int((first-start)*1000)}ms read={read}"
    except Exception as e:
        return None, f"error:{type(e).__name__}:{e}"

def bench_one(reg: str, tries: int, max_bytes: int, timeout: float):
    url = normalize_registry(reg) + SEARCH_PATH.lstrip("/")
    best, note = None, ""
    for _ in range(tries):
        t, n = http_read_time(url, max_bytes, timeout)
        note = n
        if t is not None:
            best = t if best is None else min(best, t)
    return reg, best, note

def rank_regs(cands, tries, max_bytes, timeout):
    seen, uniq = set(), []
    for u in cands:
        u = u.strip().rstrip("/")
        if u and u not in seen:
            uniq.append(u)
            seen.add(u)
    results = [bench_one(u, tries, max_bytes, timeout) for u in uniq]
    results.sort(key=lambda x: float("inf") if x[1] is None else x[1])
    return results

def load_text(path):
    try:
        with open(path, "r", encoding="utf-8") as f:
            return f.read()
    except FileNotFoundError:
        return ""
    except Exception as e:
        raise e

def save_with_backup(path, content):
    if os.path.exists(path):
        ts = datetime.now().strftime("%Y%m%d-%H%M%S")
        backup = f"{path}.bak-{ts}"
        shutil.copy2(path, backup)
        bakmsg = f"已备份到 {backup}"
    else:
        bakmsg = "(首次创建，无备份)"
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        f.write(content)
    return bakmsg

def upsert_bunfig_text(orig: str, registry_url: str, scopes: dict):
    """
    以最小侵入方式更新/插入：
    [install]
    registry = "<url>"
    [install.scopes]
    "@scope" = "<url>"
    """
    reg_line = f'registry = "{registry_url}"'
    has_install = re.search(r'^\[install\]\s*$', orig, flags=re.M)
    if has_install:
        def repl_install(block):
            body = block.group(0)
            if re.search(r'^\s*registry\s*=\s*".*?"\s*$', body, flags=re.M):
                body = re.sub(r'^\s*registry\s*=\s*".*?"\s*$', reg_line, body, flags=re.M)
            else:
                body = body + ("\n" if not body.endswith("\n") else "") + reg_line + "\n"
            return body
        orig = re.sub(r'(?ms)^\[install\]\s*$(.*?)(?=^\[|\Z)', repl_install, orig)
    else:
        orig += ("" if orig.endswith("\n") or orig == "" else "\n") + "[install]\n" + reg_line + "\n"

    if scopes:
        if not re.search(r'^\[install\.scopes\]\s*$', orig, flags=re.M):
            orig += "\n[install.scopes]\n"
        for k, v in scopes.items():
            line = f'"{k}" = "{normalize_registry(v)}"'
            pattern = rf'^\s*"{re.escape(k)}"\s*=\s*.*$'
            if re.search(pattern, orig, flags=re.M):
                orig = re.sub(pattern, line, orig, flags=re.M)
            else:
                orig = re.sub(r'(?ms)^\[install\.scopes\]\s*$(.*?)(?=^\[|\Z)',
                              lambda m: m.group(0) + ("" if m.group(1).endswith("\n") else "\n") + line + "\n",
                              orig)
    return orig

def main():
    ap = argparse.ArgumentParser(description="测速并设置 Bun 默认 registry（默认写本地 bunfig.toml；--global 写入 $HOME/.bunfig.toml）")
    ap.add_argument("--apply", action="store_true", help="写入 bunfig.toml（不加则仅测速）")
    ap.add_argument("--global", dest="use_global", action="store_true",
                    help="写入全局 $HOME/.bunfig.toml（默认写本地 ./bunfig.toml）")
    ap.add_argument("--candidates", type=str, default="",
                    help="自定义候选 registry（逗号分隔）")
    ap.add_argument("--tries", type=int, default=DEFAULT_TRIES, help="每个候选的尝试次数（默认2）")
    ap.add_argument("--bytes", type=int, default=DEFAULT_BYTES, help="读取字节上限（默认128KiB）")
    ap.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT, help="超时秒数（默认8）")
    ap.add_argument("--scope", action="append", default=[],
                    help='可重复：为 scope 指定 registry，如：--scope "@myorg=https://registry.myorg.com/"')
    args = ap.parse_args()

    cands = list(DEFAULT_CANDIDATES)
    if args.candidates.strip():
        cands += [x.strip() for x in args.candidates.split(",") if x.strip()]

    print("== Bun registry 测速开始 ==")
    print(f"候选数：{len(cands)} | 读取上限：{args.bytes} B | 重试：{args.tries}")
    results = rank_regs(cands, args.tries, args.bytes, args.timeout)

    best = None
    print("\n== 速度排名（越上越快）==")
    for i, (u, t, note) in enumerate(results, 1):
        speed = "失败" if t is None else (f"{int(t*1000)} ms" if t < 1 else f"{t:.2f} s")
        print(f"{i:02d}. {u:<48} -> {speed}   [{note}]")
        if best is None and t is not None:
            best = u

    if not best:
        print("\n未找到可用 registry。请检查网络或更换候选。")
        return

    best_norm = normalize_registry(best)
    print(f"\n最快候选：{best_norm}")

    if not args.apply:
        print("\n（仅测速未写入；加 --apply 写入 bunfig.toml。默认本地，--global 为全局。）")
        return

    # 解析 scope 形如 @org=url
    scopes = {}
    for s in args.scope:
        if "=" not in s:
            raise SystemExit(f"--scope 参数格式错误：{s}（应为 @scope=url）")
        k, v = s.split("=", 1)
        scopes[k.strip()] = v.strip()

    # 目标路径：默认 local；仅 --global 时改为全局
    if args.use_global:
        home = os.path.expanduser("~")
        path = os.path.join(home, ".bunfig.toml")
    else:
        path = os.path.join(os.getcwd(), "bunfig.toml")

    orig = load_text(path)
    newtxt = upsert_bunfig_text(orig, best_norm, scopes)
    bakmsg = save_with_backup(path, newtxt)
    print(f"\n已写入：{path}\n{bakmsg}")
    print("\n建议执行：bun install 以验证。")

if __name__ == "__main__":
    main()
